from typing import Dict, List, Optional, Tuple
import json
import re
from llm.model_factory import ModelFactory
from .authorization_service import AuthorizationService
from config import MODEL_NAME

class SQLPermissionService:
    def __init__(self):
        """初始化 SQL 权限服务"""
        self.model_name = MODEL_NAME
        self.model_client = ModelFactory.create_model()
        self.auth_service = AuthorizationService()
        
    def get_table_privileges(self, tables: List[str], user_id: str) -> Dict[str, Dict]:
        """
        从权限表中获取用户对指定表的权限
        
        Args:
            tables: 表名列表
            user_id: 用户ID
            
        Returns:
            Dict: 权限上下文字典
        """
        privileges, error = self.auth_service.get_table_privileges(tables, user_id)
        print(f"privileges: {privileges}")

        tables = privileges.keys()
        if error:
            print(f"Warning: {error}")
            # 如果发生错误，返回最严格的权限（不允许访问任何列）
            return {
                table: {
                    "allowed_columns": [],
                    "row_filter": []
                }
                for table in tables
            }
        #初始化
        privilege_context = {}

        for table in tables:
            table_privileges = privileges.get(table, {})
            if table_privileges:
                privilege_context[table] = {
                    "allowed_columns": table_privileges.get("allowed_columns", []),
                    "row_filter": table_privileges.get("row_filter", [])
                }
            else:
                # 如果表没有权限数据，返回默认的最严格权限
                privilege_context[table] = {
                    "allowed_columns": [],
                    "row_filter": []
                }

        print("privilege_context:", privilege_context)
        return privilege_context
    
        
    def extract_tables_from_sql(self, sql: str) -> List[str]:
        """
        从 SQL 语句中提取表名
        
        Args:
            sql: SQL 语句
            
        Returns:
            List[str]: 表名列表
        """
        # 使用模型来解析 SQL 并提取表名
        prompt = {
            "system":"你是一个 SQL 解析专家。你的任务是从 SQL 语句中提取所有表名，并严格返回 JSON 数组格式。"
            ,
              "user": f"SQL 语句：{sql}\n请直接返回 JSON 数组，包含完整的表名路径，例如：[\"users\"]"
              }

        
        response = self.model_client.chat(
            messages=[
                {"role": "system", "content": prompt["system"]},
                {"role": "user", "content": prompt["user"]}
            ],
            model="qwen2.5:3b"
        )
        print(f"大模型返回的原始结果:",response)

        try:
            
            tables = json.loads(response)
            return tables if isinstance(tables, list) else []
        except Exception as e:
            print("Json解析失败:", e)
            return []
            
    def rewrite_sql_with_privileges(self, original_sql, user_id):

        # privilege_context = self.get_table_privileges([], user_id) 
        privilege_context, error = self.auth_service.get_table_privileges([], user_id)
        if error:
            return '', error

        #1.处理权限上下文，转化为json字符串格式
        # privilege_context_str = json.dumps(privilege_context, indent=2, ensure_ascii=False)

        #2.Create a prompt that includes both the privilege context and the original SQL query.
        prompt = f"""You are an expert SQL security auditor.

    Below is the privilege context that defines the allowed access for each table:
    {privilege_context}

    The original SQL query is:
    {original_sql}

    Please rewrite the SQL query so that:
    1. Ensure you can find the table names in the priviledge table list, if tables not in the list, please return "error: no priviledge to run the query".
    2. It selects only the allowed columns for each table.
    3. It applies any necessary row-level filters.
    4. Ensure the rewritten SQL follows correct SQL syntax and is executable.
    5. Make sure all string literals are properly quoted (e.g., 'value' not value).
    6. Avoid using SQL keywords as column aliases in the rewritten query.
    7. Return only the final SQL query, without any explanations, comments, or markdown formatting.
    
    /no_think
    """
        #3.调用LLM进行SQL重写
        try:
            response = self.model_client.chat(
                messages=[
                    {"role": "system", "content": "You are a helpful SQL rewriting assistant."},
                    {"role": "user", "content": prompt}
                ],
                model=self.model_name,
                temperature=0
            )
            
            #4.解析LLM响应，获取重写SQL
            rewritten_sql = self.extract_sql(response)
            print("rewritten_sql: ", rewritten_sql)

            if rewritten_sql.startswith("error"):
                return '', rewritten_sql
            
            return rewritten_sql, ''
        except Exception as e:
            error_msg = f"重写SQL过程发生错误: {str(e)}"
            print(error_msg)
            return '', error_msg

    def extract_sql(self, llm_response: str) -> str:
        """
        Example:
        ```python
        vn.extract_sql("Here's the SQL query in a code block: ```sql\nSELECT * FROM customers\n```")
        ```

        Extracts the SQL query from the LLM response. This is useful in case the LLM response contains other information besides the SQL query.
        Override this function if your LLM responses need custom extraction logic.

        Args:
            llm_response (str): The LLM response.

        Returns:
            str: The extracted SQL query.
        """

        # If the llm_response contains a CTE (with clause), extract the last sql between WITH and ;
        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL)
        if sqls:
            sql = sqls[-1]
            return sql

        # If the llm_response is not markdown formatted, extract last sql by finding select and ; in the response
        sqls = re.findall(r"SELECT.*?;", llm_response, re.DOTALL)
        if sqls:
            sql = sqls[-1]
            return sql

        # If the llm_response contains a markdown code block, with or without the sql tag, extract the last sql from it
        sqls = re.findall(r"```sql\n(.*)```", llm_response, re.DOTALL)
        if sqls:
            sql = sqls[-1]
            return sql

        sqls = re.findall(r"```(.*)```", llm_response, re.DOTALL)
        if sqls:
            sql = sqls[-1]
            return sql

        return llm_response